'''
测试数据集精度统计
对比两者标注
# labelmepose 和 pose_test_dataset100-labelme ()

labelmebox 和 labelmebox
'''

import glob
import json
import os
import random
import time

import cv2
import numpy as np


# from log import get_logger
from loguru import logger  # 日志

def filter_annotation(true_annotation): # 过滤策略

    # polygon 2 rectangle
    for true_rect_shape_dict in true_annotation['shapes']:

        if true_rect_shape_dict['shape_type'] == 'polygon':  # 转矩形
            ps = true_rect_shape_dict['points']
            ps_np = np.array(ps)
            true_x1 = ps_np[:, 0].min()
            true_x2 = ps_np[:, 0].max()
            true_y1 = ps_np[:, 1].min()
            true_y2 = ps_np[:, 1].max()
            true_rect_shape_dict['points'] = [[true_x1, true_y1], [true_x2, true_y2]]
            true_rect_shape_dict['shape_type'] = 'rectangle'

    # # 删除重复标签
    # for true_rect_shape_dict in true_annotation['shapes']:

    # 删除不在clsnames里的标签
    return true_annotation

def compare_annotation(test_config):
    def iou(box1, box2):
        """
        计算两个矩形框的 IOU

        参数:
        box1 -- 第一个矩形框 [x1, y1, x2, y2]
        box2 -- 第二个矩形框 [x1, y1, x2, y2]

        返回:
        iou -- IOU 比值 (float)
        """

        # 获取坐标
        x1_1, y1_1, x2_1, y2_1 = box1
        x1_2, y1_2, x2_2, y2_2 = box2

        # 计算交集区域的坐标
        inter_x1 = max(x1_1, x1_2)
        inter_y1 = max(y1_1, y1_2)
        inter_x2 = min(x2_1, x2_2)
        inter_y2 = min(y2_1, y2_2)

        # 如果没有交集
        if inter_x1 >= inter_x2 or inter_y1 >= inter_y2:
            return 0.0

        # 计算交集面积
        inter_area = (inter_x2 - inter_x1) * (inter_y2 - inter_y1)

        # 计算每个矩形框的面积
        area1 = (x2_1 - x1_1) * (y2_1 - y1_1)
        area2 = (x2_2 - x1_2) * (y2_2 - y1_2)

        # 并集面积
        union_area = area1 + area2 - inter_area

        # IOU
        iou_value = inter_area / union_area

        return iou_value

    true_annotation_dir = test_config.get('true_annotation_dir')
    pre_annotation_dir = test_config.get('pre_annotation_dir')

    a = len(glob.glob(f'{true_annotation_dir}/*.json'))
    b = len(glob.glob(f'{pre_annotation_dir}/*.json'))
    if a != b:
        print(f'tru {a}  != pre {b}')
        raise ValueError('true_annotation_dir和pre_annotation_dir的数量不一致')

    img_glob = f'{true_annotation_dir}/*.jpg'  # 测试图片



    is_save_img = test_config.get('is_save_img', False)
    save_dir = test_config.get('save_dir', None)
    iou_thre = test_config.get('iou_thre', 0.5)

    # test_cls_names = ['baggage_place', 'none_action']
    # pre_cls_names = ['IP', 'NP']

    test_cls_names = ['aircraft', 'cargo_door', 'chock', 'refueling_truck', 'lifting_platform',
                      'transfer_car', 'staff', 'connecting', 'catering_truck', 'belt_loader',
                      'power_cable', 'pre_conditioning_air', 'thumbs_up', 'ULD', 'bag', 'push_back_tractor',
                      'no_detection']
    # pre_cls_names = ['aircraft', 'cargo_door', 'chock', 'refueling_truck', 'lifting_platform',
    #                  'transfer_car', 'staff', 'connecting', 'catering_truck', 'belt_loader',
    #                  'power_cable', 'pre_conditioning_air', 'thumbs_up', 'ULD', 'bag', 'push_back_tractor',
    #                  'no_detection']

    pre_tru_clsname_dict = {
        'aircraft': 'aircraft',
        'cargo_door': 'cargo_door',
        'chock': 'chock',
        'refueling_truck': 'refueling_truck',
        'lifting_platform': 'lifting_platform',
        'transfer_car': 'transfer_car',
        'staff': 'staff',
        'connecting': 'connecting',
        'catering_truck': 'catering_truck',
        'belt_loader': 'belt_loader',
        'power_cable': 'power_cable',
        # 'pre_conditioning_air': 'ULD', # 改
        'pre_conditioning_air': 'pre_conditioning_air',  # 改
        'thumbs_up': 'thumbs_up',
        # 'ULD': 'pre_conditioning_air', # 改
        'ULD': 'ULD',  # 改
        'bag': 'bag',
        'push_back_tractor': 'push_back_tractor',
        'no_detection': 'no_detection',
    }

    cls_colors = [(255, 255, 255), (0, 0, 0), (255, 0, 0), (0, 255, 0), (0, 0, 255),
                  (0, 255, 255), (255, 255, 0), (255, 0, 255), (255, 255, 255), (0, 0, 0),
                  (255, 0, 0), (0, 255, 0), (0, 0, 0), (255, 0, 0), (0, 255, 0),
                  (0, 255, 0), (0, 255, 0), ]

    tru_pre_acc_matrix = {}
    for tru_name in pre_tru_clsname_dict.values():
        tru_pre_acc_matrix[tru_name] = {}
        for pre_name in pre_tru_clsname_dict.keys():
            tru_pre_acc_matrix[tru_name][pre_name] = 0

    box_error_count = 0
    img_ls = glob.glob(img_glob)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    for img_path in img_ls:
        img_file_name = os.path.basename(img_path)  # '000000.jpg'
        img_name = os.path.splitext(img_file_name)[0]  # '000000'
        img = cv2.imread(img_path)
        pre_annotation_path = f"{pre_annotation_dir}/{img_name}.json"  # 根据文件名匹配
        true_annotation_path = f"{true_annotation_dir}/{img_name}.json"
        # if not os.path.exists(pre_annotation_path) or not os.path.exists(true_annotation_path): # 导致遍历img_ls和pre_annotation有关
        #     continue

        if not os.path.exists(true_annotation_path):
            true_annotation = {'shapes': []}
        else:
            true_annotation = json.load(open(true_annotation_path))
        if not os.path.exists(pre_annotation_path):
            pre_annotation = {'shapes': []}
        else:
            pre_annotation = json.load(open(pre_annotation_path))

        # 先过滤
        true_annotation = filter_annotation(true_annotation)
        pre_annotation = filter_annotation(pre_annotation)

        for true_rect_shape_dict in true_annotation['shapes']:  # 以true_annotation优先
            if not true_rect_shape_dict['shape_type'] in ['rectangle', 'polygon']:
                continue
            if true_rect_shape_dict['shape_type'] == 'polygon':  # 转矩形
                ps = true_rect_shape_dict['points']
                ps_np = np.array(ps)
                true_x1 = ps_np[:, 0].min()
                true_x2 = ps_np[:, 0].max()
                true_y1 = ps_np[:, 1].min()
                true_y2 = ps_np[:, 1].max()
            else:
                [true_x1, true_y1], [true_x2, true_y2] = true_rect_shape_dict['points']
            true_cls_name = true_rect_shape_dict['label']

            if true_cls_name not in test_cls_names: # 不考虑不在test_cls_names的
                continue
            true_x1, true_y1, true_x2, true_y2 = int(true_x1), int(true_y1), int(true_x2), int(true_y2)
            # 画 true 标注
            true_color = cls_colors[test_cls_names.index(true_cls_name)]
            cv2.rectangle(img, (true_x1, true_y1), (true_x2, true_y2), true_color, 2)  # bgr 白
            cv2.line(img, (true_x1, (true_y1 + true_y2) // 2), (true_x2, (true_y1 + true_y2) // 2), true_color, 2)  # 横线

            # 找匹配的 pre_rect_shape_dict
            max_iou_value = 0
            max_iou_pre_rect_shape_dict_ind = 0
            for shape_dict_ind, pre_rect_shape_dict in enumerate(pre_annotation['shapes']):
                if not pre_rect_shape_dict['shape_type'] in ['rectangle', 'polygon']:
                    continue
                if pre_rect_shape_dict['shape_type'] == 'polygon':  # 转矩形
                    ps = pre_rect_shape_dict['points']
                    ps_np = np.array(ps)
                    pre_x1 = ps_np[:, 0].min()
                    pre_x2 = ps_np[:, 0].max()
                    pre_y1 = ps_np[:, 1].min()
                    pre_y2 = ps_np[:, 1].max()
                else:
                    [pre_x1, pre_y1], [pre_x2, pre_y2] = pre_rect_shape_dict['points']
                # [pre_x1, pre_y1], [pre_x2, pre_y2] = pre_rect_shape_dict['points']
                iou_value = iou([true_x1, true_y1, true_x2, true_y2], [pre_x1, pre_y1, pre_x2, pre_y2])

                if iou_value > max_iou_value:
                    max_iou_value = iou_value
                    max_iou_pre_rect_shape_dict_ind = shape_dict_ind

            is_error_box = max_iou_value < iou_thre
            if is_error_box:  # 漏检
                box_error_count += 1
                # print(f'loujian: {img_name}')
                # continue
                pre_cls_name = 'no_detection'
                pre_x1, pre_y1, pre_x2, pre_y2 = 0, 0, 0, 0
            else:
                max_iou_pre_rect_shape_dict = pre_annotation['shapes'][max_iou_pre_rect_shape_dict_ind]
                pre_cls_name = max_iou_pre_rect_shape_dict['label']
                if pre_cls_name not in test_cls_names:
                    pre_cls_name = 'no_detection'
                [pre_x1, pre_y1], [pre_x2, pre_y2] = max_iou_pre_rect_shape_dict['points']
                pre_x1, pre_y1, pre_x2, pre_y2 = int(pre_x1), int(pre_y1), int(pre_x2), int(pre_y2)

            tru_pre_acc_matrix[true_cls_name][pre_cls_name] += 1 # 更新混淆矩阵

            # 画结果图
            pre_color = cls_colors[test_cls_names.index(pre_tru_clsname_dict[pre_cls_name])]
            cv2.rectangle(img, (pre_x1, pre_y1), (pre_x2, pre_y2), pre_color, 6, lineType=cv2.LINE_4)  # 黑
            cv2.line(img, ((pre_x1 + pre_x2) // 2, pre_y1), ((pre_x1 + pre_x2) // 2, pre_y2), pre_color, 6)

            text_color = (0, 255, 0) if true_cls_name == pre_tru_clsname_dict[pre_cls_name] else (0, 0, 255)
            cv2.putText(img, f'{true_cls_name}/{pre_cls_name} {max_iou_value:.2f}', (true_x1, true_y1),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2)
            if is_save_img:
                cv2.imwrite(f'{save_dir}/{img_file_name}', img, [int(cv2.IMWRITE_JPEG_QUALITY), 30])

    # print(f'img_glob {img_glob}')
    # print(f'pre_annotation_dir {pre_annotation_dir}')
    # print(f'true_annotation_dir {true_annotation_dir}')
    # print(f'save_dir {save_dir}')
    # # print(f'cls_model_type {cls_model_type}')
    # print(f'box_error_count {box_error_count}')
    avg_acc = 0
    lose_cls_count = 0
    for tru_name in test_cls_names:  # 遍历每个类别
        pre_dict = tru_pre_acc_matrix[tru_name]
        acc = pre_dict[tru_name] / max(1, sum(list(pre_dict.values())))
        print(f'{tru_name} \t{pre_dict[tru_name]} / {sum(list(pre_dict.values()))}  \t{acc * 100:.2f}%')
        if sum(list(pre_dict.values())) == 0: #
            lose_cls_count += 1
        else:
            avg_acc += acc
    avg_acc = avg_acc / (len(test_cls_names) - lose_cls_count)
    print(f'ALL {avg_acc * 100:.2f}%')


def test_compare_annotation():
    # test_config = {
    #     'true_annotation_dir': r'D:\DATA\20250611HKBZ\caitu\testdata170',  # 含img json
    #     'pre_annotation_dir': r'D:\DATA\20250611HKBZ\temp\20250615221132_yolov8mSeg_416_allv25\annotated',  # 含img json
    #     'save_dir': r'D:\DATA\20250611HKBZ\temp\20250614172530_yolov8mSeg_416_allv24\val_show_img',  # v8seg res
    #     'is_save_img': True,
    # }

    root_save_dir = r'D:\DATA\20250611HKBZ\temp\20250618141520_detr-r101-dc5_800_traindata580'

    test_config = {
        'true_annotation_dir': r'D:\DATA\20250611HKBZ\caitu\testdata170',  # 含img json
        'pre_annotation_dir': rf'{root_save_dir}\annotated',  # 含img json
        'save_dir': rf'{root_save_dir}\val_show_img',  # v8seg res
        'is_save_img': True,
        'iou_thre': 0.1
    }
    compare_annotation(test_config)


def test_compare_annotation_getpr():  # 基于compare_annotation

    # root_save_dir = r'D:\DATA\20250611HKBZ\temp\20250618141520_detr-r101-dc5_800_traindata580'
    # root_save_dir = r'D:\DATA\20250611HKBZ\temp\20250620101413_yolov8mSeg_640_allv22'
    root_save_dir = r'D:\DATA\20250611HKBZ\temp\20250620103642_YoloV8mSeg_train640_DBclip_ep200_infer1280'
    test_config = {
        'true_annotation_dir': r'D:\DATA\20250611HKBZ\caitu\before_20250619\testdata170',  # 含img json
        'pre_annotation_dir': rf'{root_save_dir}\annotated',  # 含img json
        'save_dir': rf'{root_save_dir}\val_show_img',  # v8seg res
        'is_save_img': True,
        'iou_thre': 0.1,
        # 'test_clsnames': ['aircraft', 'cargo_door', 'chock', 'refueling_truck', 'lifting_platform',
        #               'transfer_car', 'staff', 'connecting', 'catering_truck', 'belt_loader',
        #               'power_cable', 'pre_conditioning_air', 'thumbs_up', 'ULD', 'bag', 'push_back_tractor',
        #               'no_detection']
    }
    print('\nRecall pre/tru...\n')
    compare_annotation(test_config)

    true_annotation_dir = test_config['true_annotation_dir']
    test_config['true_annotation_dir'] = test_config['pre_annotation_dir']
    test_config['pre_annotation_dir'] = true_annotation_dir  # 调换顺序
    test_config['save_dir'] = test_config['save_dir'] + '_p'
    print('\nPrecision tru/pre...\n')
    compare_annotation(test_config)

if __name__ == '__main__':
    # #
    test_compare_annotation_getpr()
    # test_compare_annotation()


